{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import json\n",
    "import os\n",
    "import argparse\n",
    "from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction\n",
    "import nltk\n",
    "import numpy as np\n",
    "\n",
    "def load_json_file(file_path):\n",
    "    \"\"\"Load JSON file and return as Python object\"\"\"\n",
    "    try:\n",
    "        with open(file_path, 'r', encoding='utf-8') as f:\n",
    "            return json.load(f)\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading {file_path}: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "def extract_video_name(path):\n",
    "    \"\"\"Extract video filename without extension from path\"\"\"\n",
    "    return os.path.splitext(os.path.basename(path))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_path = 'output/batch_video_results.json'\n",
    "references_path = 'datasets/video_prompts.json'\n",
    "\n",
    "# Load JSON files\n",
    "predictions_data = load_json_file(prediction_path)\n",
    "references_data = load_json_file(references_path)\n",
    "\n",
    "if not predictions_data or not references_data:\n",
    "    print(\"Failed to load one or both JSON files. Exiting.\")\n",
    "\n",
    "# Create a dictionary of reference prompts keyed by video name\n",
    "reference_dict = {}\n",
    "for item in references_data:\n",
    "    # Adapt this according to your reference JSON structure\n",
    "    if isinstance(item, dict) and 'video' in item and 'conversations' in item:\n",
    "        reference_dict[item['video']] = item['conversations'][1]['value']\n",
    "\n",
    "\n",
    "# Create a list to store all comparison data\n",
    "comparison_data = []\n",
    "\n",
    "# Process each prediction\n",
    "for pred in predictions_data:\n",
    "    video_name = pred['video_name']\n",
    "    prediction = pred.get('prediction', '')\n",
    "    \n",
    "    # Find matching reference\n",
    "    reference = reference_dict.get(video_name, '')\n",
    "    \n",
    "    \n",
    "    # Add to comparison data\n",
    "    comparison_data.append({\n",
    "        'video_name': video_name,\n",
    "        'reference': reference,\n",
    "        'prediction': prediction,\n",
    "    })\n",
    "\n",
    "# Convert to DataFrame\n",
    "df = pd.DataFrame(comparison_data)\n",
    "#Save to parquet\n",
    "df.to_parquet('output/compare_evalutaion_results.parquet', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import re\n",
    "\n",
    "result_df = pd.read_parquet('output/compare_evalutaion_results.parquet')\n",
    "\n",
    "\n",
    "def clean_sentence(s):\n",
    "    s = s.replace(\"，\", \",\")  \n",
    "    s = re.sub(r\"-[a-z]+\\s*\\d+\", \"\", s)  \n",
    "    s = s.strip()\n",
    "    return s\n",
    "\n",
    "predictions = [clean_sentence(p) for p in result_df['prediction']]\n",
    "references = [[clean_sentence(r)] for r in result_df['reference']]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation For Video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bert-Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "compare_promts = pd.read_parquet('XXX')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert_score import score\n",
    "\n",
    "def bert_sim_eval(candidate_sentence, target_sentence):\n",
    "    # Put sentences in lists since bert-score's score function expects list inputs\n",
    "    candidates = [candidate_sentence]\n",
    "    references = [target_sentence]\n",
    "\n",
    "    # Calculate BertScore\n",
    "    P, R, F1 = score(candidates, references, lang='en', verbose=True, model_type='bert-base-uncased')\n",
    "    return {\"P\":P.item(),\"R\":R.item(),\"F1\":F1.item()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "P_ours = []\n",
    "R_ours = []\n",
    "F1_ours = []\n",
    "for index,item in tqdm(compare_promts.iterrows(),total=len(compare_promts)):    \n",
    "    res = bert_sim_eval(item['reference'],item['prediction'])\n",
    "    P_ours.append(res['P'])\n",
    "    R_ours.append(res['R'])\n",
    "    F1_ours.append(res['F1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sum(P_ours)/len(P_ours))\n",
    "print(sum(R_ours)/len(R_ours))\n",
    "print(sum(F1_ours)/len(F1_ours))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## xclip-Video Similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import CogVideoXPipeline\n",
    "from diffusers.utils import export_to_video\n",
    "import gc\n",
    "\n",
    "\n",
    "def gen_pic(compare_df):\n",
    "    pipe = CogVideoXPipeline.from_pretrained(\n",
    "        \"THUDM/CogVideoX-2b\",\n",
    "        torch_dtype=torch.float16\n",
    "    ).to(\"cuda\")\n",
    "\n",
    "    # pipe.enable_model_cpu_offload()\n",
    "    # pipe.enable_sequential_cpu_offload()\n",
    "    pipe.vae.enable_slicing()\n",
    "    pipe.vae.enable_tiling()\n",
    "    \n",
    "    for idx,row in compare_df.iterrows():\n",
    "        prompt = row['RL_generation']\n",
    "        video_name = row['video_name']\n",
    "\n",
    "        print(f\"[{idx+1}/{len(compare_df)}] Generating: RL_gen_{video_name}\")\n",
    "\n",
    "        try:\n",
    "            video = pipe(\n",
    "                prompt=prompt,\n",
    "                num_videos_per_prompt=1,\n",
    "                num_inference_steps=50,\n",
    "                num_frames=49,\n",
    "                guidance_scale=6,\n",
    "                generator=torch.Generator(device=\"cuda\").manual_seed(42),\n",
    "            ).frames[0]\n",
    "\n",
    "            export_to_video(video, f\"output/RL_genvideos/RL_gen_{video_name}\", fps=8)\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"❌ Error generating {video_name}: {e}\")\n",
    "\n",
    "        finally:\n",
    "            # Explicitly release GPU memory\n",
    "            del video\n",
    "            torch.cuda.empty_cache()\n",
    "            gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_pic(compare_promts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import av\n",
    "import torch\n",
    "import numpy as np\n",
    "from transformers import AutoProcessor, AutoModel\n",
    "from huggingface_hub import hf_hub_download\n",
    "from torch.nn.functional import cosine_similarity\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "\n",
    "def read_video_pyav(container, indices):\n",
    "    frames = []\n",
    "    container.seek(0)\n",
    "    start_index = indices[0]\n",
    "    end_index = indices[-1]\n",
    "    for i, frame in enumerate(container.decode(video=0)):\n",
    "        if i > end_index:\n",
    "            break\n",
    "        if i >= start_index and i in indices:\n",
    "            frames.append(frame)\n",
    "    return np.stack([x.to_ndarray(format=\"rgb24\") for x in frames])\n",
    "\n",
    "def sample_frame_indices(clip_len, frame_sample_rate, seg_len):\n",
    "    converted_len = int(clip_len * frame_sample_rate)\n",
    "    end_idx = np.random.randint(converted_len, seg_len)\n",
    "    start_idx = end_idx - converted_len\n",
    "    indices = np.linspace(start_idx, end_idx, num=clip_len)\n",
    "    indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)\n",
    "    return indices\n",
    "\n",
    "def get_video_features(video_path, clip_len=8):\n",
    "    \"\"\"Extract features from a video file using XCLIP model\"\"\"\n",
    "    container = av.open(video_path)\n",
    "    indices = sample_frame_indices(clip_len=clip_len, frame_sample_rate=1, \n",
    "                                 seg_len=container.streams.video[0].frames)\n",
    "    video = read_video_pyav(container, indices)\n",
    "    \n",
    "    processor = AutoProcessor.from_pretrained(\"microsoft/xclip-base-patch32\")\n",
    "    model = AutoModel.from_pretrained(\"microsoft/xclip-base-patch32\")\n",
    "    \n",
    "    inputs = processor(videos=list(video), return_tensors=\"pt\")\n",
    "    video_features = model.get_video_features(**inputs)\n",
    "    return video_features\n",
    "\n",
    "def compute_video_similarity(video1_path, video2_path):\n",
    "    \"\"\"Compute cosine similarity between two videos\"\"\"\n",
    "    features1 = get_video_features(video1_path)\n",
    "    features2 = get_video_features(video2_path)\n",
    "    \n",
    "    # Compute cosine similarity between video features\n",
    "    similarity = cosine_similarity(features1, features2)\n",
    "    return similarity.item()\n",
    "\n",
    "def compare_videos(compare_df):\n",
    "    \"\"\"Compare original videos with generated videos\"\"\"\n",
    "    results = []\n",
    "    similarity_scores = []\n",
    "\n",
    "    for idx, row in tqdm(compare_df.iterrows(), total=len(compare_df)):\n",
    "        original_video = row['video_name']\n",
    "        generated_video = f\"RL_gen_{original_video}\"\n",
    "        \n",
    "        try:\n",
    "            similarity = compute_video_similarity(\n",
    "                f\"vidprom/cog_videos_example/{original_video}\",\n",
    "                f\"output/genvideos/{generated_video}\"\n",
    "            )\n",
    "            \n",
    "            # Save intermediate results to JSON file\n",
    "            result = {\n",
    "                'video_name': original_video,\n",
    "                'generated_video': generated_video,\n",
    "                'reference': row['reference'],\n",
    "                'origin': row['Origin_generation'],\n",
    "                'DI_prediction': row['DI_generation'],\n",
    "                'video_similarity_score': similarity,\n",
    "                'promt_similarity_score': row['similarity_score']\n",
    "            }\n",
    "            \n",
    "            # Append to results list\n",
    "            results.append(result)\n",
    "            similarity_scores.append(similarity)\n",
    "\n",
    "            # Save current results to JSON file\n",
    "            result_path = 'output/RL_gen_video_comparison_results.json'\n",
    "            with open(result_path, 'w') as f:\n",
    "                json.dump(results, f, indent=4)\n",
    "                \n",
    "            print(f\"Processed {idx+1}/{len(compare_df)} videos. Current similarity: {similarity:.4f}\")\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Error comparing videos {original_video}: {e}\")\n",
    "            continue\n",
    "    \n",
    "    # Create DataFrame with results\n",
    "    results_df = pd.DataFrame(results)\n",
    "    results_df = results_df.sort_values('similarity_score', ascending=False)\n",
    "    \n",
    "    return result_path,results_df,similarity_scores\n",
    "\n",
    "# Example usage:\n",
    "# results = compare_videos(compare_df)\n",
    "# print(results.head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_path,results,video_cos_sim_scores = compare_videos(compare_promts)\n",
    "print(results.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "vid_sim_res = pd.read_json(result_path)\n",
    "print(vid_sim_res.shape)\n",
    "vid_sim_res.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cogvid",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
